#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <spu_intrinsics.h>
#include <spu_mfcio.h>

#include "fish.h"
#include "flocking.h"

#define BUCKETS 2

/* Flocking Parameters */
const double C0 = 0.00000001; // weight of vector away from neighbors within crowding radius
//const double C1 = 2000.0; // weight of vector towards center of world
const double C1 = 0.0;
const double C2 = 1.0; // weight of velocity correcting vectors towards avg. flocking velocity
const double C3 = 0.0; // weight of vector towards center of gravity
const double C4 = 0.05; // weight of random vector
const double C5 = 0.1; // weight of vector pointing to cog of local flockmates
const double R_CROWD2 = 0.00390625; // square of crowding radius
const double R_SIGHT2 = 0.015625; // square of sighting radius
const double R_SIGHT = 0.125; // sighting radius
const double BIRD_MASS = 1.0; // bird mass
const double R = 1.0; // repulsion constant
const double M = 1.0; // constant that converts vectors to acceleration
const double START_CONST = 0.5;
const double SIGHT_ANGLE = 200.0 * M_PI / 180.0;
const double SCALE_FACTOR = 1.0;
double v0x, v0y, v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y, v5x, v5y, vx, vy, ax, ay; // vectors
double cogx, cogy; // center of gravity coordinates
double local_cogx, local_cogy;
double tx, ty; // translation vectors for the moving grid

fish_t* fish[BUCKETS][BUCKETS];
int fish_bucket_n[BUCKETS][BUCKETS];
int n_fish;

void interact_fish_bucket(int, int);

int main(unsigned long long spe, unsigned long long argp, unsigned long long envp)
{
  // DMA in the spe_params_t struct
  spe_params_t spe_params __attribute__((aligned(16)));
  spe_auxinfo_t spe_auxinfo __attribute__((aligned(16)));
  int tag;

  // issue the transfer
  tag = 1;
  spu_mfcdma64(&spe_params, mfc_ea2h(argp), mfc_ea2l(argp), sizeof(spe_params_t), tag, MFC_GET_CMD);
  
  tag = 2;
  spu_mfcdma64(&spe_auxinfo, mfc_ea2h(envp), mfc_ea2l(envp), sizeof(spe_auxinfo_t), tag, MFC_GET_CMD);

  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);
  spu_writech(MFC_WrTagMask, 1 << 2);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);

  unsigned int seed = spe_auxinfo.random_seed*((unsigned int)(-1));

  srand48(seed);

  // allocate space for the incomming fish
  // n_fish is the sum of fish in each quadrant
  n_fish = spe_params.nw_n + spe_params.ne_n + spe_params.sw_n + spe_params.se_n;
  fish_bucket_n[0][0] = spe_params.nw_n;
  fish_bucket_n[0][1] = spe_params.ne_n;
  fish_bucket_n[1][0] = spe_params.sw_n;
  fish_bucket_n[1][1] = spe_params.se_n;

  fish[0][0] = (fish_t*)memalign(16, n_fish*sizeof(fish_t));
  fish[0][1] = (fish_t*)memalign(16, n_fish*sizeof(fish_t));
  fish[1][0] = (fish_t*)memalign(16, n_fish*sizeof(fish_t));
  fish[1][1] = (fish_t*)memalign(16, n_fish*sizeof(fish_t));

  // aiight, now we just have to go get the fish!!!!
  tag = 1;
  spu_mfcdma64(fish[0][0], mfc_ea2h(spe_params.nw), mfc_ea2l(spe_params.nw), 
	       spe_params.nw_n*sizeof(fish_t), tag, MFC_GET_CMD);
  tag = 2;
  spu_mfcdma64(fish[0][1], mfc_ea2h(spe_params.ne), mfc_ea2l(spe_params.ne), 
	       spe_params.ne_n*sizeof(fish_t), tag, MFC_GET_CMD);
  tag = 3;
  spu_mfcdma64(fish[1][0], mfc_ea2h(spe_params.sw), mfc_ea2l(spe_params.sw), 
	       spe_params.sw_n*sizeof(fish_t), tag, MFC_GET_CMD);
  tag = 4;
  spu_mfcdma64(fish[1][1], mfc_ea2h(spe_params.se), mfc_ea2l(spe_params.se), 
	       spe_params.se_n*sizeof(fish_t), tag, MFC_GET_CMD);

  // I'm not sure if these are required to finish in order... so we'll have 
  // to wait on all of them

  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);
  spu_writech(MFC_WrTagMask, 1 << 2);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);
  spu_writech(MFC_WrTagMask, 1 << 3);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);
  spu_writech(MFC_WrTagMask, 1 << 4);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);

  // by this time we have the fish array filled up, and the fish_buckets_n 
  // array filled up... now we can run the interact_fish_bucket function from
  // our old code just fine.

  int n = spe_auxinfo.bucket_number/BUCKETS;
  int m = spe_auxinfo.bucket_number%BUCKETS;
  unsigned long long writeback_ea;
  int writeback_num;

  interact_fish_bucket( n, m );
  
  if( n == 0 )
    if( m == 0 ) {
      writeback_ea = spe_params.nw;
      writeback_num = spe_params.nw_n;
    }
    else {
      writeback_ea = spe_params.ne;
      writeback_num = spe_params.ne_n;
    }
  else
    if( m == 0 ) {
      writeback_ea = spe_params.sw;
      writeback_num = spe_params.sw_n;
    }
    else {
      writeback_ea = spe_params.se;
      writeback_num = spe_params.se_n;
    }


  // now we need to write the fish back to main memory
  tag = 1;
  spu_mfcdma64(fish[n][m], mfc_ea2h(writeback_ea), mfc_ea2l(writeback_ea), 
	       writeback_num*sizeof(fish_t), tag, MFC_PUT_CMD);
  

  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);
  
  return 0;
}


void
interact_fish_bucket(int n, int m)
{
  int i, j, ni, mi, n0, m0;
  fish_t *__restrict__ wfish;
  fish_t const *__restrict__ rfish;
  double r2, r16, r32, r4, r, a, angle, vtempx, vtempy, dx, dy;
  int in_sight, in_crowd;

  wfish = (fish_t *__restrict__)fish[n][m];

  for(i = 0; i < fish_bucket_n[n][m]; i++){

    v0x = v0y = v1x = v1y = v2x = v2y = v3x = v3y = v4x = v4y = v5x = v5y = vx = vy = ax = ay = 0.0;
    in_sight = in_crowd = 0;
    local_cogx = 0.0;
    local_cogy = 0.0;

    // calculate random unit vector
    v4x = drand48() - 0.5;
    v4y = drand48() - 0.5;
    r2 = f_max(v4x * v4x + v4y * v4y, min_r2);
    r = sqrt(r2);
    v4x /= r;
    v4y /= r;

    // interact with fish in this bucket and the 3 neighbor buckets
    for (ni = 0; ni < 2; ni++) {
      for (mi = 0; mi < 2; mi++) {

        n0 = (n+ni+BUCKETS) % BUCKETS;
        m0 = (m+mi+BUCKETS) % BUCKETS;
        rfish = (fish_t const *__restrict__)fish[n0][m0];

        for(j = 0; j < fish_bucket_n[n0][m0]; j++){

          //dx = fmod((rfish[j].x - wfish[i].x), (2.0 - R_SIGHT));
          //dy = fmod((rfish[j].y - wfish[i].y), (2.0 - R_SIGHT));
          dx = rfish[j].x - wfish[i].x;
          if ( fabs( dx ) > WALL_SEP - R_SIGHT ) {
            dx = fabs( dx ) - WALL_SEP;
          }
          dy = rfish[j].y - wfish[i].y;
          if ( fabs( dy ) > WALL_SEP - R_SIGHT ) {
            dy = fabs( dy ) - WALL_SEP;
          }

          if (dx != 0 || dy != 0) { 
            r2 = f_max(dx * dx + dy * dy, min_r2);

            // figure out if the boid is in this boid's range of vision
            angle = fabs( atan2(dy, dx) - atan2(wfish[i].vy, wfish[i].vx) );

            if( (angle <= (SIGHT_ANGLE/2.0)) || (angle >= (2*M_PI - SIGHT_ANGLE/2.0)) ) {
	      
              // calculate crowding vector
              if ( r2 < R_CROWD2 ) {
                v0x -= dx;
                v0y -= dy;
                in_crowd++;
              }

              // calculate flocking vector
              if( r2 < R_SIGHT2 ) {
                v2x += rfish[j].vx;
                v2y += rfish[j].vy;

                local_cogx += rfish[j].x;
                local_cogy += rfish[j].y;

                in_sight++;
              }
            }
          }
        }   
      }
    }

    if( in_sight != 0 ) {
      // calculate vector pointing in the direction of the average velocity of nearby flockmates
      v2x /= in_sight;
      v2y /= in_sight;

      v2x = v2x - wfish[i].vx;
      v2y = v2y - wfish[i].vy;

      local_cogx /= in_sight;
      local_cogy /= in_sight;
      
      // calculate vector towards cog of local flockmates
      v5x = local_cogx - wfish[i].x;
      v5y = local_cogy - wfish[i].y;
    }

    if( in_crowd != 0 ) {
      v0x /= in_crowd;
      v0y /= in_crowd;
      
      r2 = v0x * v0x + v0y * v0y;

      // let the repulsion decrease with r4
      v0x /= r2*r2;
      v0y /= r2*r2;
    }

    // now we have v0 to v5...
    vx = C0*v0x + C1*v1x + C2*v2x + C3*v3x + C4*v4x + C5*v5x;
    vy = C0*v0y + C1*v1y + C2*v2y + C3*v3y + C4*v4y + C5*v5y;   


    ax = M*vx;
    ay = M*vy;
    

    wfish[i].ax = ax;
    wfish[i].ay = ay;
  }
}


